// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/dwio/dwrf/writer/writer_sink.h>

namespace kumo::pollux::dwrf {
void WriterSink::addBuffer(dwio::common::DataBuffer<char> buffer) {
  const auto length = buffer.size();
  if (length > 0) {
    if (shouldChecksum()) {
      checksum_->update(buffer.data(), length);
    }
    if (shouldCache()) {
      if (getCurrentCacheSize() + length > maxCacheSize_) {
        exceedsLimit_ = true;
        truncateCache();
      } else {
        // capture input to the cache
        auto remainingSize = length;
        auto* src = buffer.data();
        auto capacity = cacheBuffer_.capacity();
        auto size = cacheBuffer_.size();
        // resize the buffer assuming we need to use all of it
        cacheBuffer_.resize(capacity);
        while (remainingSize > 0) {
          POLLUX_CHECK_LT(size, capacity);
          const auto toWriteSize = std::min(capacity - size, remainingSize);
          std::memcpy(cacheBuffer_.data() + size, src, toWriteSize);
          size += toWriteSize;
          if (size == capacity) {
            cacheHolder_.take(cacheBuffer_);
            size = 0;
          }
          remainingSize -= toWriteSize;
          src += toWriteSize;
        }
        cacheBuffer_.resize(size);
      }
    }
  }
  if (shouldBuffer_) {
    buffers_.push_back(std::move(buffer));
    size_ += length;
  } else {
    sink_->write(std::move(buffer));
  }
}

void WriterSink::init(memory::MemoryPool& pool) {
  POLLUX_CHECK(!initialized_);
  POLLUX_CHECK(offsets_.empty());
  initialized_ = true;

  if (cacheMode_ != StripeCacheMode::NA) {
    offsets_.push_back(0);
    cacheBuffer_.reserve(SLICE_SIZE);
  }

  // initialize the buffer with the orc header
  addBuffer(pool, ORC_MAGIC.data(), ORC_MAGIC_LEN);
}
} // namespace kumo::pollux::dwrf
